import onnxruntime as ort
import numpy as np
import cv2
from PIL import Image
from torch import nn
import torch

def sigmoid(x):
    return 1/1(1+np.exp(-x))

def normalize_image(image, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):
    """
    对输入的图像数据进行标准化处理。

    参数:
    image (numpy.ndarray): 输入的图像数据，形状可以是 (H, W, C) 或 (C, H, W)。
    mean (list): 每个通道的均值，默认值为 [0.485, 0.456, 0.406]。
    std (list): 每个通道的标准差，默认值为 [0.229, 0.224, 0.225]。

    返回:
    numpy.ndarray: 标准化后的图像数据。
    """
    # 检查图像数据的形状
    if image.ndim == 3 and image.shape[2] == 3:  # (H, W, C) 格式
        image = np.transpose(image, (2, 0, 1))  # 转换为 (C, H, W) 格式

    # 确保输入的图像数据的像素值范围在 [0, 1] 之间
    image = image.astype(np.float32) / 255.0

    # 进行标准化处理
    mean = np.array(mean).reshape(-1, 1, 1)
    std = np.array(std).reshape(-1, 1, 1)
    normalized_image = (image - mean) / std

    # 如果之前进行了维度调整，将数据转换回 (H, W, C) 格式
    if image.ndim == 3 and image.shape[0] == 3:
        normalized_image = np.transpose(normalized_image, (1, 2, 0))

    return normalized_image




class HumanAttributeOnnx(nn.Module):
    def __init__(self, onnx_model_path='/Users/luotianhang/Documents/VScode-work/floatingAI/attr_0816_64_color_skip_quick_36/onnx/pedestrian_attribute#38p#250227.onnx'):
        super().__init__()

        self.session = ort.InferenceSession(onnx_model_path)

        self.input_name  = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name
        self.input_shape = self.session.get_inputs()[0].shape

        print('input_name:', self.input_name)
        print('input_shape:',self.input_shape)
        print('output_name:',self.output_name)



    def forward(self,x,use_sigmoid=True):
        x = x.cpu().numpy()
        outputs = self.session.run([self.output_name], {self.input_name: x})
        outputs = torch.tensor(outputs[0])
        if use_sigmoid:
            outputs = torch.sigmoid(outputs)

        return outputs
